package org.jetlinks.pro.cassandra.metadata;

import io.netty.buffer.ByteBufUtil;
import io.netty.buffer.Unpooled;
import lombok.AllArgsConstructor;
import lombok.Getter;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.codec.binary.Base64;
import org.hswebframework.ezorm.rdb.executor.SqlRequest;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.ByteBuffer;
import java.util.function.Function;

@Getter
@AllArgsConstructor(staticName = "of")
@Slf4j
public class CassandraPagingContext {

    private int pageIndex;

    private int pageSize;

    private ByteBuffer pageState;

    @Setter
    private ByteBuffer nextPageState;

    public static CassandraPagingContext of(int pageSize) {
        return of(0, pageSize, null, null);
    }

    public static CassandraPagingContext of(int pageIndex, int pageSize, String base64) {
        return of(pageIndex, pageSize, StringUtils.isEmpty(base64) ? null : ByteBuffer.wrap(Base64.decodeBase64(base64)), null);
    }

    public CassandraPagingContext nextPage() {
        return CassandraPagingContext.of(pageIndex + 1, pageSize, nextPageState, null);
    }

    int getRealPageSize() {
        if (pageIndex != 0 && pageState == null) {
            //设置了分页索引,但是没有设置分页状态,说明客户端不支持游标分页
            //则查询全部数据然后只获取对应页多数据
            return (pageIndex + 1) * pageSize;
        }
        return pageSize;
    }

    int getFetchSkip() {
        if (pageIndex != 0 && pageState == null) {
            return pageIndex * pageSize;
        }
        return 0;
    }

    public String nextPageStateBase64() {
        if (nextPageState == null) {
            return null;
        }
        return Base64.encodeBase64String(ByteBufUtil.getBytes(Unpooled.wrappedBuffer(nextPageState)));
    }

    public static Mono<CassandraPagingContext> current() {
        return Mono
            .subscriberContext()
            .handle((ctx, sink) -> {
                if (ctx.hasKey(CassandraPagingContext.class)) {
                    sink.next(ctx.get(CassandraPagingContext.class));
                }
            });
    }

    public static Mono<CassandraPagingContext> current(SqlRequest sqlRequest) {
        if (sqlRequest instanceof PagingSqlRequest) {
            PagingSqlRequest pagingSqlRequest = ((PagingSqlRequest) sqlRequest);

            return current()
                .doOnNext(ctx -> {
                    ctx.pageSize = pagingSqlRequest.getPageSize();
                    ctx.pageState = pagingSqlRequest.getPageState();
                    ctx.pageIndex = pagingSqlRequest.getPageIndex();
                })
                .defaultIfEmpty(CassandraPagingContext.of(pagingSqlRequest.getPageIndex(),
                                                          pagingSqlRequest.getPageSize(),
                                                          pagingSqlRequest.getPageState(), null));
        }
        return current();
    }

    public static <T> Function<Mono<T>, Mono<T>> setMonoPaging(CassandraPagingContext context) {
        return mono -> mono.subscriberContext(ctx -> ctx.put(CassandraPagingContext.class, context));
    }

    public static <T> Function<Flux<T>, Flux<T>> setFluxPaging(CassandraPagingContext context) {
        return flux -> flux.subscriberContext(ctx -> ctx.put(CassandraPagingContext.class, context));
    }

}
